#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File          :    train.py    
@Contact       :    LJL959@QQ.COM
@License       :    (C)Copyright 2020-2021, Liugroup-NLPR-CASIA
@Modify Time   :    2020/12/9 13:45        
@Author        :    LiuJiaolong
@Version       :    1.0
@Description    :   训练网络代码
"""
# import lib
import torch
import numpy as np
import torch.nn as nn
from D2L.MyFirstDemoWithFMRIData.network import Net

data_dir = 'D:\Study\PycharmProjects\D2L\MyFirstDemoWithFMRIData\data/trainData'
lr = 0.0001     # 学习率
n_epoch = 10
batch_size = 3  # batch_size大小
model_dir = 'D:\Study\PycharmProjects\D2L\MyFirstDemoWithFMRIData\model/'
all_train_data =[]


def train():

    data_with_label = np.load(data_dir + '/drug_health_data.npy')
    health_data = data_with_label[: 20]     # 健康人为前20个数据
    drug_data = data_with_label[20:]        # 病人为后20个数据
    for i in range(5):
        np.random.shuffle(health_data)
        np.random.shuffle(drug_data)
    train_health_data, train_drug_data = health_data[:17], drug_data[:17]   # 训练数据为0~17个
    test_health_data, test_drug_data = health_data[17:], drug_data[17:]     # 测试数据为17~20个

    train_data = np.vstack((train_health_data, train_drug_data))            # 训练数据拼接
    test_data = np.vstack((test_health_data, test_drug_data))               # 测试数据拼接

    model = Net()           # 实例化一个网络
    model = model           # 使用CPU训练，如果有GPU的话后面加.cuda()
    # model = nn.DataParallel(model)      # 多GPU并行化处理网络
    model.train()           # 训练网络
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)         # 实例化一个优化器，调整网络参数，优化方式为Adam方法
    criterion = torch.nn.CrossEntropyLoss()                         # 定义loss计算方法，交叉熵，可以理解为两者数值越接近其值越小
    cnt = 0                 # 训练数据数量

    for epoch in range(n_epoch):                                    # 读取数据集进行训练
        for batch_idx, data in enumerate(data_with_label):          # 从封装好的npy文件中读取数据与对应标签
            label_temp = data_with_label[batch_idx, 116, 0]         # 标签只占一个位置,在第116行 第0列位置处
            label_x = torch.tensor(label_temp, dtype=torch.long)    # 数据类型转换为tensor

            # 从封装好的文件里取得数据有问题
            data_temp = data_with_label[batch_idx, 0:116, 0:150]    # 前0~116存放结点数据，0~150存放特征，所以数据就是0~150全取
            data_x = torch.tensor(data_temp)                        # 转换成tensor
            data_x_unsqueeze = torch.unsqueeze(data_x, 0)           # 扩充维度
            data_new = torch.unsqueeze(data_x_unsqueeze, 0)         # 扩充维度

            # 老方法取的数据没问题
            single_data = data[:116]
            single_data = torch.Tensor(single_data)
            all_train_data.append(single_data)
            temp_data = torch.unsqueeze(all_train_data[0], 0)       # 在第一维上扩一维
            new_data = torch.unsqueeze(temp_data, 0)

            out = model(new_data)               # 计算网络输出值，就是输入网络的一个特征数据，输出病人和健康人的概率，调用了网络中的forward方法
            loss = criterion(out, label_x.unsqueeze(0))     # 计算损失，就是网络输出值和实际label的差异，差异越小，说明拟合效果越好
            loss.backward()                     # 误差反向传播，采用求导的方式，计算网络中每个结点参数的梯度，梯度越大越不合理，需要调整
            optimizer.step()                    # 优化采用设定的优化方法对网络中的各个参数进行调整
            optimizer.zero_grad()               # 清除优化器中的梯度以便下一次计算，因为优化器不会默认保留，不清除的话每次计算都会保留
            cnt += 1

            print('Epoch:{0}, Frame:{1}, train_loss:{2}'.format(epoch, cnt*batch_size, loss/batch_size))    # 打印每个batch_size的训练结果

    torch.save(model.state_dict(), '{0}/model.pth'.format(model_dir))


if __name__ == '__main__':
    train()





